import os
import sys
import yaml
import pandas as pd
import numpy as np
import json
from tqdm import tqdm

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from src.data_prep.data_loader import load_ticket_data, process_ticket_features, split_train_test
from src.features.feature_engineer import TicketFeatureEngineer
from src.models.model_trainer import train_ticket_model


def generate_mock_ticket_data(num_samples: int = 1000) -> pd.DataFrame:
    """
    生成模拟工单数据用于演示
    """
    print(f"正在生成{num_samples}条模拟工单数据...")
    
    np.random.seed(42)
    data = []
    
    # 定义问题类型和对应的关键词分布
    problem_types = {
        '硬件_屏幕': {'刷新慢': 0.8, '黑屏': 0.9, '闪烁': 0.7, '电池': 0.2, '联网': 0.1},
        '硬件_电池': {'电池': 0.9, '续航短': 0.8, '充电': 0.7, '黑屏': 0.5, '发热': 0.6},
        '软件_网络': {'联网': 0.9, '刷新慢': 0.5, '超时': 0.8, '连接': 0.7, '闪退': 0.3},
        '软件_系统': {'闪退': 0.9, '卡顿': 0.8, '崩溃': 0.7, '刷新慢': 0.6, '发热': 0.4},
        '其他': {'操作': 0.5, '设置': 0.4, '功能': 0.6, '使用': 0.7}
    }
    
    # 客户类型
    customer_segments = ['超市', '便利店', '药店', '书店', '服装专卖店']
    
    # 设备年龄段
    device_age_buckets = ['0-3个月', '3-6个月', '6-12个月', '1-2年', '2年以上']
    
    for i in tqdm(range(num_samples)):
        ticket_id = f"T-{np.random.randint(2023, 2025)}{np.random.randint(1, 13):02d}-{str(i).zfill(5)}"
        
        # 随机选择问题类型，但控制比例
        problem_weights = [0.3, 0.25, 0.2, 0.15, 0.1]  # 各类问题的占比权重
        problem_type = np.random.choice(list(problem_types.keys()), p=problem_weights)
        
        # 生成关键词
        keywords = {}
        base_keywords = problem_types[problem_type]
        
        # 添加基础关键词
        for keyword, base_weight in base_keywords.items():
            # 添加噪声
            noise = np.random.normal(0, 0.1)
            weight = max(0.1, min(0.99, base_weight + noise))
            keywords[keyword] = round(weight, 2)
        
        # 添加一些随机关键词
        all_keywords = ['刷新慢', '黑屏', '电池', '联网', '闪退', '发热', '闪烁', 
                       '续航短', '充电', '超时', '连接', '卡顿', '崩溃', '操作', '设置']
        random_keywords = np.random.choice(all_keywords, size=2, replace=False)
        for keyword in random_keywords:
            if keyword not in keywords:
                keywords[keyword] = round(np.random.uniform(0.1, 0.4), 2)
        
        # 客户类型
        customer_segment = np.random.choice(customer_segments)
        
        # 设备年龄段
        age_weights = [0.2, 0.25, 0.25, 0.15, 0.15]  # 设备年龄分布
        device_age_bucket = np.random.choice(device_age_buckets, p=age_weights)
        
        # 对于老旧设备，增加硬件问题概率
        if device_age_bucket in ['1-2年', '2年以上'] and problem_type == '其他':
            # 30%的概率将'其他'类型改为硬件问题
            if np.random.random() < 0.3:
                problem_type = np.random.choice(['硬件_屏幕', '硬件_电池'], p=[0.4, 0.6])
                # 更新关键词
                base_keywords = problem_types[problem_type]
                for keyword, base_weight in base_keywords.items():
                    if keyword in keywords:
                        keywords[keyword] = max(keywords[keyword], round(base_weight * 0.8 + np.random.normal(0, 0.1), 2))
                    else:
                        keywords[keyword] = round(base_weight * 0.8 + np.random.normal(0, 0.1), 2)
        
        # 将关键词转换为JSON字符串
        keywords_json = json.dumps(keywords, ensure_ascii=False)
        
        data.append({
            'ticket_id': ticket_id,
            'text_keywords': keywords_json,
            'customer_segment': customer_segment,
            'device_age_bucket': device_age_bucket,
            'problem_type': problem_type
        })
    
    return pd.DataFrame(data)


def main():
    """
    主函数
    """
    # 加载配置
    config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
                             'config', 'config.yaml')
    with open(config_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    
    # 创建数据目录
    os.makedirs(config['data']['sample_data_path'], exist_ok=True)
    
    # 生成模拟数据
    ticket_data_path = os.path.join(config['data']['sample_data_path'], 'ticket_sample_data.csv')
    if not os.path.exists(ticket_data_path):
        df = generate_mock_ticket_data(2000)
        df.to_csv(ticket_data_path, index=False)
        print(f"模拟数据已保存至: {ticket_data_path}")
    else:
        print(f"使用已有的模拟数据: {ticket_data_path}")
        df = pd.read_csv(ticket_data_path)
    
    # 处理数据
    print("开始处理数据...")
    df = process_ticket_features(df)
    
    # 特征工程
    print("开始特征工程...")
    feature_engineer = TicketFeatureEngineer()
    df = feature_engineer.build_features(df)
    
    # 准备特征和标签
    exclude_cols = ['ticket_id', 'problem_type']
    feature_cols = [col for col in df.columns if col not in exclude_cols]
    
    X = df[feature_cols]
    y = df['problem_type']
    
    # 检查数据分布
    print(f"\n数据分布:")
    for problem_type in y.unique():
        count = (y == problem_type).sum()
        print(f"{problem_type}: {count} ({count / len(y):.2%})")
    
    # 划分训练集和测试集
    print("\n划分训练集和测试集...")
    X_train, X_test, y_train, y_test = split_train_test(X, y)
    
    print(f"训练集大小: {len(X_train)}, 测试集大小: {len(X_test)}")
    
    # 训练模型
    print("开始训练工单分类模型...")
    model = train_ticket_model(
        config=config,
        X_train=X_train,
        y_train=y_train,
        X_test=X_test,
        y_test=y_test
    )
    
    print("\n模型训练完成！")
    print(f"模型保存路径: {config['models']['ticket_model_path']}")


if __name__ == "__main__":
    main()
